#include <dot/astdot.h>
#include <common.h>

AstDot::AstDot(std::ostream & out) {
  this->id = 1;
  this->dot = new Dot(out);
}

AstDot::~AstDot() {
  if(this->dot != nullptr)
    delete this->dot;
}

void AstDot::visit(ast::NodeCompUnit & comp_unit) {
  std::string str;
  comp_unit.dot_string(str);
  int id_buff = this->id;
  add_declare(id_buff, str);
  for(auto & global : comp_unit.global_list) {
    add_direction(id_buff, this->id+1);
    this->id++;
    if(auto func_def =
        std::get_if<std::unique_ptr<ast::NodeFuncDef>>(&global)) {
      (*func_def)->dot_accept(*this);
    } else if(auto global_decl =
        std::get_if<std::unique_ptr<ast::NodeGlobalDecl>>(&global)) {
      (*global_decl)->dot_accept(*this);
    } else {
      Panic("undefined option");
    }
  }
}

void AstDot::visit(ast::NodeGlobalDecl & global_decl) {
  std::string str;
  global_decl.dot_string(str);
  add_declare(this->id, str);
  add_direction(this->id, this->id+1);
  this->id++;
  global_decl.global_defs->dot_accept(*this);
}

void AstDot::visit(ast::NodeGlobalDefs & global_defs) {
  std::string str;
  global_defs.dot_string(str);
  int id_buff = this->id;
  add_declare(id_buff, str);
  for(const auto & global_def : global_defs.global_def_list) {
    add_direction(id_buff, this->id+1);
    this->id++;
    global_def->dot_accept(*this);
  }
}

void AstDot::visit(ast::NodeGlobalDef & global_def) {
  std::string str;
  global_def.dot_string(str);
  add_declare(this->id, str);
}

void AstDot::visit(ast::NodeConstDecl & const_decl) {
  std::string str;
  const_decl.dot_string(str);
  add_declare(this->id, str);
  add_direction(this->id, this->id+1);
  this->id++;
  const_decl.const_defs->dot_accept(*this);
}

void AstDot::visit(ast::NodeConstDefs & const_defs) {
  std::string str;
  const_defs.dot_string(str);
  int id_buff = this->id;
  add_declare(id_buff, str);
  for(auto & const_def : const_defs.const_def_list) {
    add_direction(id_buff, this->id+1);
    this->id++;
    const_def->dot_accept(*this);
  }
}

void AstDot::visit(ast::NodeConstDef & const_def) {
  std::string str;
  const_def.dot_string(str);
  add_declare(this->id, str);
}

void AstDot::visit(ast::NodeVarDecl & var_decl) {
  std::string str;
  var_decl.dot_string(str);
  add_declare(this->id, str);
  add_direction(this->id, id+1);
  this->id++;
  var_decl.var_defs->dot_accept(*this);
}

void AstDot::visit(ast::NodeVarDefs & var_defs) {
  std::string str;
  var_defs.dot_string(str);
  add_declare(this->id, str);
  int id_buff = this->id;
  for(const auto & var_def : var_defs.var_def_list) {
    this->id++;
    add_direction(id_buff, this->id);
    var_def->dot_accept(*this);
  }
}

void AstDot::visit(ast::NodeVarDef & var_def) {
  std::string str;
  var_def.dot_string(str);
  add_declare(this->id, str);
  int id_buff = this->id;
  switch(var_def.stage) {
  case ast::NodeVarDef::Stage::EXP_INIT:
    add_direction(id_buff, this->id+1);
    this->id++;
    var_def.exp->dot_accept(*this);
    break;
  case ast::NodeVarDef::Stage::EXP_NOT_INIT:
    // empty
    break;
  case ast::NodeVarDef::Stage::ARR_ZERO_INIT:
    add_direction(id_buff, this->id+1);
    this->id++;
    var_def.array_const->dot_accept(*this);
    break;
  case ast::NodeVarDef::Stage::ARR_INIT:
    add_direction(id_buff, this->id+1);
    this->id++;
    var_def.array_const->dot_accept(*this);
    add_direction(id_buff, this->id+1);
    this->id++;
    var_def.array_init->dot_accept(*this);
    break;
  case ast::NodeVarDef::Stage::ARR_NOT_INIT:
    add_direction(id_buff, this->id+1);
    this->id++;
    var_def.array_const->dot_accept(*this);
    break;
  default: Panic("var def stage undefined");
  }
}

void AstDot::visit(ast::NodeArray & array) {
  std::string str;
  array.dot_string(str);
  add_declare(this->id, str);
  for(auto & exp : array.exp_list) {
    add_direction(this->id, this->id+1);
    this->id++;
    exp->dot_accept(*this);
  }
}

void AstDot::visit(ast::NodeArrayConst & array_const) {
  std::string str;
  array_const.dot_string(str);
  add_declare(this->id, str);
}

void AstDot::visit(ast::NodeArrayInit & array_init) {
  std::string str;
  array_init.dot_string(str);
  add_declare(this->id, str);
  int id_buff = this->id;
  for(auto & init : array_init.init_list) {
    add_direction(id_buff, this->id+1);
    this->id++;
    if(auto exp = std::get_if<std::unique_ptr<ast::NodeExp>>(&init)) {
      (*exp)->dot_accept(*this);
    }else if(auto array_init = std::get_if<std::unique_ptr<ast::NodeArrayInit>>(&init)) {
      (*array_init)->dot_accept(*this);
    }
  }
}

void AstDot::visit(ast::NodeFuncDef & func_def) {
  std::string str;
  func_def.dot_string(str);
  int id_buff = this->id;
  add_declare(id_buff, str);
  if(func_def.stage == ast::NodeFuncDef::HAVE_ARGS) {
    add_direction(id_buff, this->id+1);
    this->id++;
    func_def.func_def_args->dot_accept(*this);
  }
  add_direction(id_buff, this->id+1);
  this->id++;
  func_def.block->dot_accept(*this);
}

void AstDot::visit(ast::NodeFuncDefArgs & func_def_args) {
  std::string str;
  func_def_args.dot_string(str);
  add_declare(this->id, str);
}

void AstDot::visit(ast::NodeFuncCall & func_call) {
  std::string str;
  func_call.dot_string(str);
  add_declare(this->id, str);
  if(func_call.func_args.has_value()) {
    add_direction(this->id, this->id+1);
    this->id++;
    func_call.func_args.value()->dot_accept(*this);
  }
}

void AstDot::visit(ast::NodeBlock & block) {
  std::string str;
  block.dot_string(str);
  add_declare(this->id, str);
  add_direction(this->id, this->id+1);
  this->id++;
  block.block_item_list->dot_accept(*this);
}

void AstDot::visit(ast::NodeBlockItems & block_items) {
  std::string str;
  block_items.dot_string(str);
  add_declare(this->id, str);
  int id_buff = id;
  for(const auto & stmt : block_items.stmt_list) {
    this->id++;
    add_direction(id_buff, this->id);
    stmt->dot_accept(*this);
  }
  for(const auto & var_decl : block_items.var_decl_list) {
    this->id++;
    add_direction(id_buff, this->id);
    var_decl->dot_accept(*this);
  }
  for(const auto & const_decl : block_items.const_decl_list) {
    this->id++;
    add_direction(id_buff, this->id);
    const_decl->dot_accept(*this);
  }
}

void AstDot::visit(ast::NodeStmt & stmt) {
  std::string str;
  stmt.dot_string(str);
  add_declare(this->id, str);
  add_direction(this->id, this->id+1);
  this->id++;
  switch(stmt.stage) {
  case ast::NodeStmt::Stage::ASSIGN:
    stmt.stmt_assign->dot_accept(*this);
    break;
  case ast::NodeStmt::Stage::RETURN:
    stmt.stmt_return->dot_accept(*this);
    break;
  case ast::NodeStmt::Stage::IF:
    stmt.stmt_if->dot_accept(*this);
    break;
  case ast::NodeStmt::Stage::WHILE:
    stmt.stmt_while->dot_accept(*this);
    break;
  case ast::NodeStmt::Stage::BREAK:
    stmt.stmt_break->dot_accept(*this);
    break;
  case ast::NodeStmt::Stage::CONTINUE:
    stmt.stmt_continue->dot_accept(*this);
    break;
  case ast::NodeStmt::Stage::FUNCCALL:
    stmt.stmt_func_call->dot_accept(*this);
    break;
  case ast::NodeStmt::Stage::BLOCK:
    stmt.block->dot_accept(*this);
    break;
  default:
    Panic("stmt stage undefined");
  }
}

void AstDot::visit(ast::NodeStmtAssign & stmt_assign) {
  std::string str;
  stmt_assign.dot_string(str);
  add_declare(this->id, str);
  add_direction(this->id, this->id+1);
  this->id++;
  stmt_assign.exp->dot_accept(*this);
}

void AstDot::visit(ast::NodeStmtReturn & stmt_return) {
  std::string str;
  stmt_return.dot_string(str);
  add_declare(this->id, str);
  if(stmt_return.exp.has_value()) {
    add_direction(this->id, this->id+1);
    this->id++;
    stmt_return.exp.value()->dot_accept(*this);
  }
}

void AstDot::visit(ast::NodeStmtIf & stmt_if) {
  std::string str;
  stmt_if.dot_string(str);
  add_declare(this->id, str);
  int buff_id = this->id;
  switch(stmt_if.stage) {
  case ast::NodeStmtIf::Stage::NO_ELSE:
    add_direction(buff_id, this->id+1, "true");
    this->id++;
    stmt_if.stmt_true->dot_accept(*this);
    break;
  case ast::NodeStmtIf::Stage::ELSE:
    add_direction(buff_id, this->id+1, "true");
    this->id++;
    stmt_if.stmt_true->dot_accept(*this);
    add_direction(buff_id, this->id+1, "false");
    this->id++;
    stmt_if.stmt_false->dot_accept(*this);
    break;
  default: Panic("stmt if stage undefined");
  }
}

void AstDot::visit(ast::NodeStmtWhile & stmt_while) {
  std::string str;
  stmt_while.dot_string(str);
  add_declare(this->id, str);
  int buff_id = this->id;
  add_direction(buff_id, this->id+1);
  this->id++;
  stmt_while.cond->dot_accept(*this);
  add_direction(buff_id, this->id+1);
  this->id++;
  stmt_while.stmt->dot_accept(*this);
}

void AstDot::visit(ast::NodeStmtBreak & stmt_break) {
  std::string str;
  stmt_break.dot_string(str);
  add_declare(this->id, str);
}

void AstDot::visit(ast::NodeStmtContinue & stmt_continue) {
  std::string str;
  stmt_continue.dot_string(str);
  add_declare(this->id, str);
}

void AstDot::visit(ast::NodeStmtFuncCall & stmt_func_call) {
  std::string str;
  stmt_func_call.dot_string(str);
  add_declare(this->id, str);
  add_direction(this->id, this->id+1);
  this->id++;
  stmt_func_call.func_call->dot_accept(*this);
}

void AstDot::visit(ast::NodeExp & exp) {
  std::string str;
  exp.dot_string(str);
  add_declare(this->id, str);
  int buff_id = this->id;
  switch(exp.stage) {
  case ast::NodeExp::EXP:
    add_direction(buff_id, this->id+1);
    this->id++;
    exp.exp1->dot_accept(*this);
    add_direction(buff_id, this->id+1);
    this->id++;
    exp.exp2->dot_accept(*this);
    break;
  case ast::NodeExp::BRAC_EXP:
    add_direction(this->id, this->id+1);
    this->id++;
    exp.exp1->dot_accept(*this);
    break;
  case ast::NodeExp::INT_NUMBER:
    str += R"(\nint_number: )" + std::to_string(exp.int_number);
    break;
  case ast::NodeExp::FLOAT_NUMBER:
    str += R"(\nfloat_number: )" + std::to_string(exp.float_number);
    break;
  case ast::NodeExp::IDEN:
    str += R"(\niden: )" + exp.iden;
    break;
  case ast::NodeExp::ARRAY:
    add_direction(this->id, this->id+1);
    this->id++;
    exp.array->dot_accept(*this);
    break;
  case ast::NodeExp::FUNC:
    exp.func_call->dot_accept(*this);
    break;
  default: Panic("exp.stage undefined");
  }
}

void AstDot::visit(ast::NodeFuncArgs & func_args) {
  std::string str;
  func_args.dot_string(str);
  int id_buff = this->id;
  add_declare(id_buff, str);
  for(auto & exp : func_args.exp_list) {
    add_direction(id_buff, this->id+1);
    this->id++;
    exp->dot_accept(*this);
  }
}

void AstDot::visit(ast::NodeCond & cond) {
  std::string str;
  cond.dot_string(str);
  add_declare(this->id, str);
  add_direction(this->id, this->id+1);
  this->id++;
  cond.exp->dot_accept(*this);
}

void AstDot::write() {
  for(auto declare : this->declare_list) {
    this->dot->write(std::to_string(GET_KEY(declare)) + " [label=\"" + GET_VAL(declare) + "\"]");
  }
  for(auto direction : this->direction_list) {
    this->dot->write(std::to_string(GET_KEY(direction)) + " -> "
        + std::to_string(std::get<0>(GET_VAL(direction)))
        + " [label=\"" + std::get<1>(GET_VAL(direction)) + "\"]");
  }

  delete this->dot;
  this->dot = nullptr;
}

void AstDot::add_declare(int id, std::string & label) {
  this->declare_list.insert({id, label});
}

void AstDot::add_direction(int id_src, int id_dst) {
  this->direction_list.insert({id_src, std::tuple(id_dst, "")});
}

void AstDot::add_direction(int id_src, int id_dst, std::string label) {
  this->direction_list.insert({id_src, std::tuple(id_dst, label)});
}
